import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import pandas as pd
import random
import time

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False

# --------------------------- 常量参数定义 ---------------------------
# 无人机FY1参数
FY1_INIT_POS = np.array([17800.0, 0.0, 1800.0])  # 初始位置(x,y,z)，单位：m
FY1_SPEED_MIN = 70.0                             # 最小速度，单位：m/s
FY1_SPEED_MAX = 140.0                            # 最大速度，单位：m/s

# 导弹M1参数
M1_INIT_POS = np.array([20000.0, 0.0, 2000.0])   # 初始位置(x,y,z)，单位：m
M1_SPEED = 300.0                                 # 飞行速度，单位：m/s
M1_TARGET = np.array([0.0, 0.0, 0.0])            # 假目标位置(原点)

# 真目标参数
TRUE_TARGET_BOTTOM = np.array([0.0, 200.0, 0.0]) # 下底面中心
TRUE_TARGET_RADIUS = 7.0                         # 半径，单位：m
TRUE_TARGET_HEIGHT = 10.0                        # 高度，单位：m

# 烟幕弹参数
SMOKE_RADIUS = 10.0                              # 烟幕云团半径，单位：m
SMOKE_SINK_SPEED = 3.0                           # 下沉速度，单位：m/s
SMOKE_EFFECTIVE_TIME = 20.0                      # 有效持续时间，单位：s
GRAVITY = 9.8                                    # 重力加速度，单位：m/s²

# 时间参数
M1_FLIGHT_DISTANCE = np.linalg.norm(M1_INIT_POS - M1_TARGET)
M1_TOTAL_TIME = M1_FLIGHT_DISTANCE / M1_SPEED    # 导弹总飞行时间
TIME_STEP = 0.5                                  # 时间步长，单位：s
TIME_POINTS = np.arange(0, M1_TOTAL_TIME + TIME_STEP, TIME_STEP)

# 优化算法参数
POPULATION_SIZE = 80
MAX_ITERATIONS = 200
C1 = 2.0  # 个体学习因子
C2 = 2.0  # 社会学习因子
W_START = 0.9  # 初始惯性权重
W_END = 0.4   # 最终惯性权重

# --------------------------- 核心功能函数 ---------------------------
def generate_target_samples() -> np.ndarray:
    """生成真目标的采样点，用于遮蔽判定"""
    samples = []
    # 下底面圆周点
    theta = np.linspace(0, 2*np.pi, 12, endpoint=False)
    for t in theta:
        x = TRUE_TARGET_BOTTOM[0] + TRUE_TARGET_RADIUS * np.cos(t)
        y = TRUE_TARGET_BOTTOM[1] + TRUE_TARGET_RADIUS * np.sin(t)
        z = TRUE_TARGET_BOTTOM[2]
        samples.append([x, y, z])
    
    # 上底面圆周点
    for t in theta:
        x = TRUE_TARGET_BOTTOM[0] + TRUE_TARGET_RADIUS * np.cos(t)
        y = TRUE_TARGET_BOTTOM[1] + TRUE_TARGET_RADIUS * np.sin(t)
        z = TRUE_TARGET_BOTTOM[2] + TRUE_TARGET_HEIGHT
        samples.append([x, y, z])
    
    # 中心轴上的点
    for z in np.linspace(TRUE_TARGET_BOTTOM[2], TRUE_TARGET_BOTTOM[2] + TRUE_TARGET_HEIGHT, 5):
        samples.append([TRUE_TARGET_BOTTOM[0], TRUE_TARGET_BOTTOM[1], z])
    
    return np.array(samples)

# 生成真目标采样点
TARGET_SAMPLES = generate_target_samples()

def get_missile_position(time: float) -> np.ndarray:
    """计算导弹在指定时间的位置"""
    if time >= M1_TOTAL_TIME:
        return M1_TARGET
    ratio = time / M1_TOTAL_TIME
    return M1_INIT_POS + ratio * (M1_TARGET - M1_INIT_POS)

def get_drone_position(angle: float, speed: float, time: float) -> np.ndarray:
    """计算无人机在指定时间的位置"""
    dx = speed * np.cos(angle) * time
    dy = speed * np.sin(angle) * time
    return np.array([
        FY1_INIT_POS[0] + dx,
        FY1_INIT_POS[1] + dy,
        FY1_INIT_POS[2]  # 保持高度不变
    ])

def calculate_smoke_effect(smoke_params: tuple, drone_angle: float, drone_speed: float) -> np.ndarray:
    """计算单枚烟幕弹的有效遮蔽时间"""
    drop_time, delay_time = smoke_params
    effective = np.zeros_like(TIME_POINTS, dtype=bool)
    
    # 计算投放点
    drop_pos = get_drone_position(drone_angle, drone_speed, drop_time)
    
    # 计算起爆点
    detonate_time = drop_time + delay_time
    detonate_pos = np.array([
        drop_pos[0] + drone_speed * np.cos(drone_angle) * delay_time,
        drop_pos[1] + drone_speed * np.sin(drone_angle) * delay_time,
        drop_pos[2] - 0.5 * GRAVITY * (delay_time ** 2)
    ])
    
    # 烟幕有效时间范围
    start_idx = np.argmin(np.abs(TIME_POINTS - detonate_time))
    end_idx = np.argmin(np.abs(TIME_POINTS - (detonate_time + SMOKE_EFFECTIVE_TIME)))
    end_idx = min(end_idx, len(TIME_POINTS) - 1)
    
    # 检查每个时间点的遮蔽效果
    for i in range(start_idx, end_idx + 1):
        t = TIME_POINTS[i]
        missile_pos = get_missile_position(t)
        
        # 计算当前烟幕中心位置（考虑下沉）
        smoke_age = t - detonate_time
        smoke_center = np.array([
            detonate_pos[0],
            detonate_pos[1],
            max(0, detonate_pos[2] - SMOKE_SINK_SPEED * smoke_age)
        ])
        
        # 检查是否遮蔽任何采样点
        for sample in TARGET_SAMPLES:
            # 计算导弹到采样点的线段是否与烟幕球相交
            a = missile_pos
            b = sample
            c = smoke_center
            r = SMOKE_RADIUS
            
            ab = b - a
            ac = c - a
            ab_len_sq = np.dot(ab, ab)
            
            if ab_len_sq < 1e-8:  # 两点重合
                if np.linalg.norm(ac) <= r:
                    effective[i] = True
                continue
            
            t = max(0, min(1, np.dot(ac, ab) / ab_len_sq))
            closest = a + t * ab
            dist = np.linalg.norm(closest - c)
            
            if dist <= r:
                effective[i] = True
                break
                
    return effective

def fitness_function(params: np.ndarray) -> float:
    """适应度函数：计算总有效遮蔽时间"""
    # 参数解析：无人机角度、速度、3枚烟幕弹的投放时间和延迟时间
    angle = params[0]
    speed = params[1]
    smoke1 = (params[2], params[5])
    smoke2 = (params[3], params[6])
    smoke3 = (params[4], params[7])
    
    # 约束检查与罚分
    penalty = 0
    
    # 速度约束
    if speed < FY1_SPEED_MIN:
        penalty += (FY1_SPEED_MIN - speed) * 100
    elif speed > FY1_SPEED_MAX:
        penalty += (speed - FY1_SPEED_MAX) * 100
    
    # 投放时间约束（时间递增且间隔至少1秒）
    if params[3] - params[2] < 1:
        penalty += (1 - (params[3] - params[2])) * 100
    if params[4] - params[3] < 1:
        penalty += (1 - (params[4] - params[3])) * 100
    
    # 时间非负约束
    for p in params[2:8]:
        if p < 0:
            penalty += (-p) * 100
    
    # 起爆点高度约束（不低于地面）
    for i in [5, 6, 7]:  # 延迟时间索引
        drop_time = params[i-3]
        delay = params[i]
        drop_pos = get_drone_position(angle, speed, drop_time)
        detonate_z = drop_pos[2] - 0.5 * GRAVITY * (delay **2)
        if detonate_z < 0:
            penalty += (-detonate_z) * 100
    
    # 计算各烟幕弹的有效遮蔽时间
    eff1 = calculate_smoke_effect(smoke1, angle, speed)
    eff2 = calculate_smoke_effect(smoke2, angle, speed)
    eff3 = calculate_smoke_effect(smoke3, angle, speed)
    
    # 总有效遮蔽时间（去重）
    total_eff = np.logical_or(np.logical_or(eff1, eff2), eff3)
    total_time = np.sum(total_eff) * TIME_STEP
    
    # 适应度 = 总有效时间 - 罚分
    return max(total_time - penalty, 0)

# --------------------------- 粒子群优化算法 ---------------------------
def pso_optimization() -> tuple:
    """使用粒子群优化算法寻找最优投放策略"""
    # 参数维度：8个参数（角度、速度、3个投放时间、3个延迟时间）
    dim = 8
    
    # 参数范围
    bounds = [
        [0, 2*np.pi],  # 角度
        [FY1_SPEED_MIN, FY1_SPEED_MAX],  # 速度
        [0, M1_TOTAL_TIME*0.6],  # 烟幕1投放时间
        [0, M1_TOTAL_TIME*0.7],  # 烟幕2投放时间
        [0, M1_TOTAL_TIME*0.8],  # 烟幕3投放时间
        [0, 15],  # 烟幕1延迟时间
        [0, 15],  # 烟幕2延迟时间
        [0, 15]   # 烟幕3延迟时间
    ]
    
    # 初始化粒子群
    particles = np.random.rand(POPULATION_SIZE, dim)
    for i in range(dim):
        particles[:, i] = particles[:, i] * (bounds[i][1] - bounds[i][0]) + bounds[i][0]
    
    # 确保投放时间递增
    for i in range(POPULATION_SIZE):
        particles[i, 3] = max(particles[i, 3], particles[i, 2] + 1)
        particles[i, 4] = max(particles[i, 4], particles[i, 3] + 1)
    
    # 初始化速度
    velocities = np.random.randn(POPULATION_SIZE, dim) * 0.1
    
    # 个体最优和全局最优
    pbest = particles.copy()
    pbest_fitness = np.array([fitness_function(p) for p in particles])
    gbest_idx = np.argmax(pbest_fitness)
    gbest = particles[gbest_idx].copy()
    gbest_fitness = pbest_fitness[gbest_idx]
    
    # 记录适应度变化
    fitness_history = [gbest_fitness]
    
    # 迭代优化
    for iter in range(MAX_ITERATIONS):
        # 动态调整惯性权重
        w = W_START - (W_START - W_END) * (iter / MAX_ITERATIONS)
        
        for i in range(POPULATION_SIZE):
            # 更新速度
            r1, r2 = np.random.rand(), np.random.rand()
            velocities[i] = (w * velocities[i] +
                           C1 * r1 * (pbest[i] - particles[i]) +
                           C2 * r2 * (gbest - particles[i]))
            
            # 限制速度范围
            for j in range(dim):
                max_vel = 0.1 * (bounds[j][1] - bounds[j][0])
                velocities[i, j] = np.clip(velocities[i, j], -max_vel, max_vel)
            
            # 更新位置
            particles[i] += velocities[i]
            
            # 边界检查
            for j in range(dim):
                particles[i, j] = np.clip(particles[i, j], bounds[j][0], bounds[j][1])
            
            # 确保投放时间顺序
            particles[i, 3] = max(particles[i, 3], particles[i, 2] + 1)
            particles[i, 4] = max(particles[i, 4], particles[i, 3] + 1)
            
            # 计算适应度
            current_fitness = fitness_function(particles[i])
            
            # 更新个体最优
            if current_fitness > pbest_fitness[i]:
                pbest[i] = particles[i].copy()
                pbest_fitness[i] = current_fitness
                
                # 更新全局最优
                if current_fitness > gbest_fitness:
                    gbest = particles[i].copy()
                    gbest_fitness = current_fitness
        
        # 记录历史
        fitness_history.append(gbest_fitness)
        
        # 打印进度
        if (iter + 1) % 20 == 0:
            print(f"迭代 {iter+1}/{MAX_ITERATIONS}, 最优适应度: {gbest_fitness:.2f}秒")
    
    return gbest, gbest_fitness, fitness_history

# --------------------------- 结果可视化与分析 ---------------------------
def visualize_results(best_params: np.ndarray):
    """可视化优化结果"""
    # 解析最优参数
    angle = best_params[0]
    speed = best_params[1]
    smoke1 = (best_params[2], best_params[5])
    smoke2 = (best_params[3], best_params[6])
    smoke3 = (best_params[4], best_params[7])
    
    # 创建图形
    fig = plt.figure(figsize=(15, 10))
    
    # 1. 轨迹可视化
    ax1 = fig.add_subplot(221, projection='3d')
    ax1.set_title('无人机、导弹及烟幕弹轨迹')
    ax1.set_xlabel('X坐标 (m)')
    ax1.set_ylabel('Y坐标 (m)')
    ax1.set_zlabel('Z坐标 (m)')
    
    # 绘制无人机轨迹
    drone_times = np.linspace(0, max(best_params[2], best_params[3], best_params[4]) + 5, 100)
    drone_positions = np.array([get_drone_position(angle, speed, t) for t in drone_times])
    ax1.plot(drone_positions[:, 0], drone_positions[:, 1], drone_positions[:, 2], 'b-', label='无人机轨迹')
    
    # 绘制导弹轨迹
    missile_positions = np.array([get_missile_position(t) for t in TIME_POINTS])
    ax1.plot(missile_positions[:, 0], missile_positions[:, 1], missile_positions[:, 2], 'r-', label='导弹轨迹')
    
    # 标记烟幕弹投放点
    drop_pos1 = get_drone_position(angle, speed, smoke1[0])
    drop_pos2 = get_drone_position(angle, speed, smoke2[0])
    drop_pos3 = get_drone_position(angle, speed, smoke3[0])
    ax1.scatter(drop_pos1[0], drop_pos1[1], drop_pos1[2], c='g', s=100, marker='o', label='烟幕弹投放点')
    ax1.scatter(drop_pos2[0], drop_pos2[1], drop_pos2[2], c='g', s=100, marker='o')
    ax1.scatter(drop_pos3[0], drop_pos3[1], drop_pos3[2], c='g', s=100, marker='o')
    
    # 标记真目标
    theta = np.linspace(0, 2*np.pi, 100)
    x = TRUE_TARGET_BOTTOM[0] + TRUE_TARGET_RADIUS * np.cos(theta)
    y = TRUE_TARGET_BOTTOM[1] + TRUE_TARGET_RADIUS * np.sin(theta)
    ax1.plot(x, y, np.full_like(theta, TRUE_TARGET_BOTTOM[2]), 'k-')
    ax1.plot(x, y, np.full_like(theta, TRUE_TARGET_BOTTOM[2]+TRUE_TARGET_HEIGHT), 'k-')
    for i in range(0, 100, 10):
        ax1.plot([x[i], x[i]], [y[i], y[i]], 
                [TRUE_TARGET_BOTTOM[2], TRUE_TARGET_BOTTOM[2]+TRUE_TARGET_HEIGHT], 'k-')
    
    ax1.legend()
    
    # 2. 遮蔽效果时间线
    ax2 = fig.add_subplot(222)
    ax2.set_title('烟幕遮蔽效果时间线')
    ax2.set_xlabel('时间 (s)')
    ax2.set_ylabel('遮蔽效果')
    ax2.set_ylim(-0.5, 3.5)
    ax2.set_yticks([0, 1, 2, 3])
    ax2.set_yticklabels(['无遮蔽', '烟幕1', '烟幕2', '烟幕3'])
    
    # 计算各烟幕的遮蔽效果
    eff1 = calculate_smoke_effect(smoke1, angle, speed)
    eff2 = calculate_smoke_effect(smoke2, angle, speed)
    eff3 = calculate_smoke_effect(smoke3, angle, speed)
    
    # 绘制遮蔽效果
    ax2.fill_between(TIME_POINTS, 0, 1, where=eff1, color='green', alpha=0.3)
    ax2.fill_between(TIME_POINTS, 1, 2, where=eff2, color='blue', alpha=0.3)
    ax2.fill_between(TIME_POINTS, 2, 3, where=eff3, color='purple', alpha=0.3)
    
    # 总遮蔽效果
    total_eff = np.logical_or(np.logical_or(eff1, eff2), eff3)
    ax2.fill_between(TIME_POINTS, 3, 3.5, where=total_eff, color='red', alpha=0.5)
    ax2.text(0, 3.6, '总遮蔽', verticalalignment='bottom')
    
    # 3. 适应度变化曲线
    ax3 = fig.add_subplot(212)
    ax3.set_title('优化过程适应度变化')
    ax3.set_xlabel('迭代次数')
    ax3.set_ylabel('最大适应度 (s)')
    ax3.plot(fitness_history, 'b-')
    
    plt.tight_layout()
    plt.show()

def save_results(best_params: np.ndarray, best_fitness: float):
    """保存优化结果到Excel文件"""
    # 解析最优参数
    angle = best_params[0]
    speed = best_params[1]
    smoke1 = (best_params[2], best_params[5])
    smoke2 = (best_params[3], best_params[6])
    smoke3 = (best_params[4], best_params[7])
    
    # 计算详细信息
    drone_info = {
        "飞行方向角(弧度)": angle,
        "飞行方向角(度)": np.degrees(angle),
        "飞行速度(m/s)": speed
    }
    
    smoke_info = []
    for i, (drop_time, delay) in enumerate([smoke1, smoke2, smoke3], 1):
        drop_pos = get_drone_position(angle, speed, drop_time)
        detonate_time = drop_time + delay
        detonate_pos = np.array([
            drop_pos[0] + speed * np.cos(angle) * delay,
            drop_pos[1] + speed * np.sin(angle) * delay,
            drop_pos[2] - 0.5 * GRAVITY * (delay **2)
        ])
        
        smoke_info.append({
            "烟幕弹编号": i,
            "投放时刻(s)": drop_time,
            "投放位置X(m)": drop_pos[0],
            "投放位置Y(m)": drop_pos[1],
            "投放位置Z(m)": drop_pos[2],
            "起爆延迟(s)": delay,
            "起爆时刻(s)": detonate_time,
            "起爆位置X(m)": detonate_pos[0],
            "起爆位置Y(m)": detonate_pos[1],
            "起爆位置Z(m)": detonate_pos[2]
        })
    
    # 计算遮蔽时间
    eff1 = calculate_smoke_effect(smoke1, angle, speed)
    eff2 = calculate_smoke_effect(smoke2, angle, speed)
    eff3 = calculate_smoke_effect(smoke3, angle, speed)
    total_eff = np.logical_or(np.logical_or(eff1, eff2), eff3)
    total_time = np.sum(total_eff) * TIME_STEP
    
    # 保存到Excel
    with pd.ExcelWriter('烟幕弹投放优化结果.xlsx') as writer:
        # 无人机信息
        pd.DataFrame([drone_info]).to_excel(writer, sheet_name='无人机参数', index=False)
        
        # 烟幕弹信息
        pd.DataFrame(smoke_info).to_excel(writer, sheet_name='烟幕弹参数', index=False)
        
        # 效果总结
        summary = pd.DataFrame([{
            "总有效遮蔽时间(s)": total_time,
            "导弹总飞行时间(s)": M1_TOTAL_TIME,
            "遮蔽覆盖率(%)": (total_time / M1_TOTAL_TIME) * 100
        }])
        summary.to_excel(writer, sheet_name='遮蔽效果', index=False)
    
    print(f"优化结果已保存到Excel文件，总有效遮蔽时间: {total_time:.2f}秒")

# --------------------------- 主函数 ---------------------------
if __name__ == "__main__":
    print("问题三求解：无人机烟幕弹投放策略优化")
    print("="*50)
    
    # 记录开始时间
    start_time = time.time()
    
    # 执行优化
    print("开始优化计算...")
    best_params, best_fitness, fitness_history = pso_optimization()
    
    # 输出优化时间
    end_time = time.time()
    print(f"\n优化完成，耗时: {end_time - start_time:.2f}秒")
    print(f"最优遮蔽时间: {best_fitness:.2f}秒")
    
    # 可视化结果
    visualize_results(best_params)
    
    # 保存结果
    save_results(best_params, best_fitness)
